#!/usr/bin/env python
# -*- coding: utf-8 -*-

""" By Martin Senande-Rivera
	For Spatial and temporal expansion of global wildland fire activity in response to climate change """

import os
import glob, sys
import numpy as np
import pandas as pd
from pandas import Series, DataFrame
import xarray as xr
import datetime as dt 

ruta='./WFDE5/'

## Precipitation
ds_tp_m = xr.open_dataset(ruta+'wfde5_tp_m_mean.nc')
ds_tp_a = xr.open_dataset(ruta+'wfde5_tp_a_mean.nc')
ds_tp = xr.open_dataset(ruta+'wfde5_tp_months.nc')


lons = np.linspace(-180.25,180.25,722)
lats = np.linspace(-90.25,90.25,362)
TP_m = xr.DataArray(np.zeros([12,362, 722]), 
 	dims=['month','lat', 'lon'], 
 	coords={'month':np.arange(1,13), 'lat': lats, 'lon': lons},
 	name='tp',
 	attrs={'units':'mm'})
TP_a = xr.DataArray(np.zeros([362, 722]), 
 	dims=['lat', 'lon'], 
 	coords={'lat': lats, 'lon': lons},
 	name='tp',
 	attrs={'units':'mm'})
TP = xr.DataArray(np.zeros([252,362, 722]), 
 	dims=['time','lat', 'lon'], 
 	coords={'time':ds_tp.time, 'lat': lats, 'lon': lons},
 	name='tp',
 	attrs={'units':'mm'})

TP_m.values[:,1:-1,1:-1] = ds_tp_m.tp.values[:]
TP_m.values[:,1:-1,0] = ds_tp_m.tp.values[:,:,-1]
TP_m.values[:,1:-1,-1] = ds_tp_m.tp.values[:,:,0]
TP_m.values[:,0,1:-1] = ds_tp_m.tp.values[:,0,:]
TP_m.values[:,-1,1:-1] = ds_tp_m.tp.values[:,-1,:]
TP_m.values[:,0,0] = ds_tp_m.tp.values[:,0,0]
TP_m.values[:,0,-1] = ds_tp_m.tp.values[:,0,-1]
TP_m.values[:,-1,0] = ds_tp_m.tp.values[:,-1,0]
TP_m.values[:,-1,-1] = ds_tp_m.tp.values[:,-1,-1]

TP_a.values[1:-1,1:-1] = ds_tp_a.tp.values[:]
TP_a.values[1:-1,0] = ds_tp_a.tp.values[:,-1]
TP_a.values[1:-1,-1] = ds_tp_a.tp.values[:,0]
TP_a.values[0,1:-1] = ds_tp_a.tp.values[0,:]
TP_a.values[-1,1:-1] = ds_tp_a.tp.values[-1,:]
TP_a.values[0,0] = ds_tp_a.tp.values[0,0]
TP_a.values[0,-1] = ds_tp_a.tp.values[0,-1]
TP_a.values[-1,0] = ds_tp_a.tp.values[-1,0]
TP_a.values[-1,-1] = ds_tp_a.tp.values[-1,-1]

TP.values[:,1:-1,1:-1] = ds_tp.tp.values[:]
TP.values[:,1:-1,0] = ds_tp.tp.values[:,:,-1]
TP.values[:,1:-1,-1] = ds_tp.tp.values[:,:,0]
TP.values[:,0,1:-1] = ds_tp.tp.values[:,0,:]
TP.values[:,-1,1:-1] = ds_tp.tp.values[:,-1,:]
TP.values[:,0,0] = ds_tp.tp.values[:,0,0]
TP.values[:,0,-1] = ds_tp.tp.values[:,0,-1]
TP.values[:,-1,0] = ds_tp.tp.values[:,-1,0]
TP.values[:,-1,-1] = ds_tp.tp.values[:,-1,-1]

ds_tp_a_interp = TP_a.interp(lat=np.linspace(-89.875,89.875,720), lon=np.linspace(-179.875,179.875,1440), method='linear')
ds_tp_m_interp = TP_m.interp(lat=np.linspace(-89.875,89.875,720), lon=np.linspace(-179.875,179.875,1440), method='linear')
ds_tp_interp = TP.interp(lat=np.linspace(-89.875,89.875,720), lon=np.linspace(-179.875,179.875,1440), method='linear')

## Snow
ds_tp_m = xr.open_dataset(ruta+'wfde5_sn_m_mean.nc')
ds_tp_a = xr.open_dataset(ruta+'wfde5_sn_a_mean.nc')
ds_tp = xr.open_dataset(ruta+'wfde5_sn_months.nc')

lons = np.linspace(-180.25,180.25,722)
lats = np.linspace(-90.25,90.25,362)
TP_m = xr.DataArray(np.zeros([12,362, 722]), 
 	dims=['month','lat', 'lon'], 
 	coords={'month':np.arange(1,13), 'lat': lats, 'lon': lons},
 	name='tp',
 	attrs={'units':'mm'})
TP_a = xr.DataArray(np.zeros([362, 722]), 
 	dims=['lat', 'lon'], 
 	coords={'lat': lats, 'lon': lons},
 	name='tp',
 	attrs={'units':'mm'})
TP = xr.DataArray(np.zeros([252,362, 722]), 
 	dims=['time','lat', 'lon'], 
 	coords={'time':ds_tp.time, 'lat': lats, 'lon': lons},
 	name='tp',
 	attrs={'units':'mm'})

TP_m.values[:,1:-1,1:-1] = ds_tp_m.tp.values[:]
TP_m.values[:,1:-1,0] = ds_tp_m.tp.values[:,:,-1]
TP_m.values[:,1:-1,-1] = ds_tp_m.tp.values[:,:,0]
TP_m.values[:,0,1:-1] = ds_tp_m.tp.values[:,0,:]
TP_m.values[:,-1,1:-1] = ds_tp_m.tp.values[:,-1,:]
TP_m.values[:,0,0] = ds_tp_m.tp.values[:,0,0]
TP_m.values[:,0,-1] = ds_tp_m.tp.values[:,0,-1]
TP_m.values[:,-1,0] = ds_tp_m.tp.values[:,-1,0]
TP_m.values[:,-1,-1] = ds_tp_m.tp.values[:,-1,-1]

TP_a.values[1:-1,1:-1] = ds_tp_a.tp.values[:]
TP_a.values[1:-1,0] = ds_tp_a.tp.values[:,-1]
TP_a.values[1:-1,-1] = ds_tp_a.tp.values[:,0]
TP_a.values[0,1:-1] = ds_tp_a.tp.values[0,:]
TP_a.values[-1,1:-1] = ds_tp_a.tp.values[-1,:]
TP_a.values[0,0] = ds_tp_a.tp.values[0,0]
TP_a.values[0,-1] = ds_tp_a.tp.values[0,-1]
TP_a.values[-1,0] = ds_tp_a.tp.values[-1,0]
TP_a.values[-1,-1] = ds_tp_a.tp.values[-1,-1]

TP.values[:,1:-1,1:-1] = ds_tp.tp.values[:]
TP.values[:,1:-1,0] = ds_tp.tp.values[:,:,-1]
TP.values[:,1:-1,-1] = ds_tp.tp.values[:,:,0]
TP.values[:,0,1:-1] = ds_tp.tp.values[:,0,:]
TP.values[:,-1,1:-1] = ds_tp.tp.values[:,-1,:]
TP.values[:,0,0] = ds_tp.tp.values[:,0,0]
TP.values[:,0,-1] = ds_tp.tp.values[:,0,-1]
TP.values[:,-1,0] = ds_tp.tp.values[:,-1,0]
TP.values[:,-1,-1] = ds_tp.tp.values[:,-1,-1]

ds_sn_a_interp = TP_a.interp(lat=np.linspace(-89.875,89.875,720), lon=np.linspace(-179.875,179.875,1440), method='linear')
ds_sn_m_interp = TP_m.interp(lat=np.linspace(-89.875,89.875,720), lon=np.linspace(-179.875,179.875,1440), method='linear')
ds_sn_interp = TP.interp(lat=np.linspace(-89.875,89.875,720), lon=np.linspace(-179.875,179.875,1440), method='linear')

ds_t_a_interp = ds_sn_a_interp+ds_tp_a_interp
ds_t_m_interp = ds_sn_m_interp+ds_tp_m_interp
ds_t_interp = ds_sn_interp+ds_tp_interp

ds_t_a_interp.to_netcdf(+'wfde5_tp_a_mean_interp025.nc')
ds_t_m_interp.to_netcdf(+'wfde5_tp_m_mean_interp025.nc')
ds_t_interp.to_netcdf(+'wfde5_tp_months_interp025.nc')


### Temperature
ds_t2m_m = xr.open_dataset(ruta+'wfde5_t2m_m_mean.nc')
ds_t2m_a = xr.open_dataset(ruta+'wfde5_t2m_a_mean.nc')
ds_t2m = xr.open_dataset(ruta+'wfde5_t2m_months.nc')


lons = np.linspace(-180.25,180.25,722)
lats = np.linspace(-90.25,90.25,362)
t2m_m = xr.DataArray(np.zeros([12,362, 722]), 
 	dims=['month','lat', 'lon'], 
 	coords={'month':np.arange(1,13), 'lat': lats, 'lon': lons},
 	name='t2m',
 	attrs={'units':'Celsius'})
t2m_a = xr.DataArray(np.zeros([362, 722]), 
 	dims=['lat', 'lon'], 
 	coords={'lat': lats, 'lon': lons},
 	name='t2m',
 	attrs={'units':'Celsius'})
t2m = xr.DataArray(np.zeros([252,362, 722]), 
 	dims=['time','lat', 'lon'], 
 	coords={'time':ds_tp.time, 'lat': lats, 'lon': lons},
 	name='t2m',
 	attrs={'units':'Celsius'})


t2m_m.values[:,1:-1,1:-1] = ds_t2m_m.t2m.values[:]
t2m_m.values[:,1:-1,0] = ds_t2m_m.t2m.values[:,:,-1]
t2m_m.values[:,1:-1,-1] = ds_t2m_m.t2m.values[:,:,0]
t2m_m.values[:,0,1:-1] = ds_t2m_m.t2m.values[:,0,:]
t2m_m.values[:,-1,1:-1] = ds_t2m_m.t2m.values[:,-1,:]
t2m_m.values[:,0,0] = ds_t2m_m.t2m.values[:,0,0]
t2m_m.values[:,0,-1] = ds_t2m_m.t2m.values[:,0,-1]
t2m_m.values[:,-1,0] = ds_t2m_m.t2m.values[:,-1,0]
t2m_m.values[:,-1,-1] = ds_t2m_m.t2m.values[:,-1,-1]

t2m_a.values[1:-1,1:-1] = ds_t2m_a.t2m.values[:]
t2m_a.values[1:-1,0] = ds_t2m_a.t2m.values[:,-1]
t2m_a.values[1:-1,-1] = ds_t2m_a.t2m.values[:,0]
t2m_a.values[0,1:-1] = ds_t2m_a.t2m.values[0,:]
t2m_a.values[-1,1:-1] = ds_t2m_a.t2m.values[-1,:]
t2m_a.values[0,0] = ds_t2m_a.t2m.values[0,0]
t2m_a.values[0,-1] = ds_t2m_a.t2m.values[0,-1]
t2m_a.values[-1,0] = ds_t2m_a.t2m.values[-1,0]
t2m_a.values[-1,-1] = ds_t2m_a.t2m.values[-1,-1]

t2m.values[:,1:-1,1:-1] = ds_t2m.t2m.values[:]
t2m.values[:,1:-1,0] = ds_t2m.t2m.values[:,:,-1]
t2m.values[:,1:-1,-1] = ds_t2m.t2m.values[:,:,0]
t2m.values[:,0,1:-1] = ds_t2m.t2m.values[:,0,:]
t2m.values[:,-1,1:-1] = ds_t2m.t2m.values[:,-1,:]
t2m.values[:,0,0] = ds_t2m.t2m.values[:,0,0]
t2m.values[:,0,-1] = ds_t2m.t2m.values[:,0,-1]
t2m.values[:,-1,0] = ds_t2m.t2m.values[:,-1,0]
t2m.values[:,-1,-1] = ds_t2m.t2m.values[:,-1,-1]


ds_t2m_a_interp = t2m_a.interp(lat=np.linspace(-89.875,89.875,720), lon=np.linspace(-179.875,179.875,1440), method='linear')
ds_t2m_m_interp = t2m_m.interp(lat=np.linspace(-89.875,89.875,720), lon=np.linspace(-179.875,179.875,1440), method='linear')
ds_t2m_interp = t2m.interp(lat=np.linspace(-89.875,89.875,720), lon=np.linspace(-179.875,179.875,1440), method='linear')

ds_t2m_a_interp.to_netcdf(ruta+'wfde5_t2m_a_mean_interp025.nc')
ds_t2m_m_interp.to_netcdf(ruta+'wfde5_t2m_m_mean_interp025.nc')
ds_t2m_interp.to_netcdf(ruta+'wfde5_t2m_months_interp025.nc')
